//
// Created by xiang on 2021/9/16.
//

#ifndef FASTER_LIO_IVOX3D_H
#define FASTER_LIO_IVOX3D_H

// #include <glog/logging.h>
// #include <execution>
// #include <list>
// #include <thread>
#include <unordered_map>
#include <boost/make_shared.hpp>

#include <omp.h>

#include <pcl/point_types.h>
#include <pcl/point_cloud.h>

#include "eigen_types.h"
#include "ivox3d_node.hpp"

namespace slam
{

    enum class IVoxNodeType
    {
        DEFAULT, // linear ivox
        PHC,     // phc ivox
    };

    /// traits for NodeType
    template <IVoxNodeType node_type, typename PointT, int dim>
    struct IVoxNodeTypeTraits
    {
    };

    template <typename PointT, int dim>
    struct IVoxNodeTypeTraits<IVoxNodeType::DEFAULT, PointT, dim>
    {
        using NodeType = IVoxNode<PointT, dim>;
    };

    template <typename PointT, int dim>
    struct IVoxNodeTypeTraits<IVoxNodeType::PHC, PointT, dim>
    {
        using NodeType = IVoxNodePhc<PointT, dim>;
    };

    template <int dim = 3, IVoxNodeType node_type = IVoxNodeType::DEFAULT, typename PointType = pcl::PointXYZ>
    class IVox
    {
    public:
        using KeyType = Eigen::Matrix<int, dim, 1>;
        using PtType = Eigen::Matrix<float, dim, 1>;
        using NodeType = typename IVoxNodeTypeTraits<node_type, PointType, dim>::NodeType;
        using PointVector = std::vector<PointType, Eigen::aligned_allocator<PointType>>;
        using DistPoint = typename NodeType::DistPoint;

        enum class NearbyType
        {
            CENTER, // center only
            NEARBY6,
            NEARBY18,
            NEARBY26,
        };

        struct Options
        {
            float resolution_ = 0.2;                       // ivox resolution
            float inv_resolution_ = 10.0;                  // inverse resolution
            NearbyType nearby_type_ = NearbyType::NEARBY6; // nearby range
            std::size_t capacity_ = 1100000;               // capacity
        };

        /**
         * constructor
         * @param options  ivox options
         */
        explicit IVox(Options options) : options_(options)
        {
            options_.inv_resolution_ = 1.0 / options_.resolution_;
            GenerateNearbyGrids();
        }

        /**
         * add points
         * @param points_to_add
         */
        void AddPoints(const PointVector &points_to_add);
        // zm--------------------------------------------------------------------------------------------------
        int is_in_area(Eigen::Vector3f c);
        void IVOX_get_areas(std::vector<float> vec_map_min, std::vector<float> vec_map_max, size_t num_area);
        void Mapping_AddPoints(const PointVector &map_points_to_add);
        // zm--------------------------------------------------------------------------------------------------
        void erasePoints(const PointVector &points_to_erase);

        /// get nn
        bool GetClosestPoint(const PointType &pt, PointType &closest_pt);

        /// get nn with condition
        bool GetClosestPoint(const PointType &pt, PointVector &closest_pt, int max_num = 5, double max_range = 5.0);

        /// get nn in cloud
        bool GetClosestPoint(const PointVector &cloud, PointVector &closest_cloud);

        typename pcl::PointCloud<PointType>::Ptr GetPointCloud();

        /// get number of points
        size_t NumPoints() const;

        /// get number of valid grids
        size_t NumValidGrids() const;

        /// get statistics of the points
        std::vector<float> StatGridPoints() const;

    private:
        /// generate the nearby grids according to the given options
        void GenerateNearbyGrids();

        /// position to grid
        KeyType Pos2Grid(const PtType &pt) const;

        Options options_;
        std::unordered_map<KeyType, typename std::list<std::pair<KeyType, NodeType>>::iterator, hash_vec<dim>>
            grids_map_;                                       // voxel hash map
        std::list<std::pair<KeyType, NodeType>> grids_cache_; // voxel cache
        std::vector<KeyType> nearby_grids_;                   // nearbys

        // zm--------------------------------------------------------------------------------------------------
        std::vector<float> vec_map_min_;
        std::vector<float> vec_map_max_;
        size_t num_area_;
    };

    // zm--------------------------------------------------------------------------------------------------
    template <int dim, IVoxNodeType node_type, typename PointType>
    int IVox<dim, node_type, PointType>::is_in_area(Eigen::Vector3f c)
    {
        float k = -2.36075470841859;
        float b = 50.5311325972969;
        int ff = 2;
        int n = 1;
        if (k * c[0] + b - c[1] < 5)
        {
            n = 0;
        }
        if (c[1] > -4.3 && c[1] < 20)
        {
            n = 0;
        }
        if (c[2] > 4.5)
        {
            n = 0;
        }
        if (k * c[0] + b - c[1] < 5)
        {
            n = 0;
        }
        if (c[1] > -4.3 && c[1] < 20)
        {
            n = 0;
        }
        if (c[2] > 4.5)
        {
            n = 0;
        }
        if (c[0] > -26.7 && c[0] < 11.6 && c[1] > -14.9 && c[1] < -13.9) // 第二个墙
        {
            n = 0;
        }
        if (c[0] > -31.5 && c[0] < 16.6 && c[1] > -26.1 && c[1] < -24.7) // 第三个墙
        {
            n = 0;
        }
        if (c[0] > -36.3 && c[0] < 21 && c[1] > -36.9 && c[1] < -35.9) // 第四个墙
        {
            n = 0;
        }
        if (c[0] > -41.5 && c[0] < 25.8 && c[1] > -48.1 && c[1] < -46.9) // 第五个墙
        {
            n = 0;
        }
        if (c[0] > -41.5 && c[0] < 30.4 && c[1] > -58.9 && c[1] < -57.9) // 第六个墙
        {
            n = 0;
        }
        if (c[0] > -41.5 && c[0] < 35 && c[1] > -69.9 && c[1] < -68.7) // 第七个墙
        {
            n = 0;
        }
        if (c[0] > -41.5 && c[0] < 44.6 && c[1] > -91.9 && c[1] < -90.7) // 第八个墙
        {
            n = 0;
        }
        if (n == 0)
        {
            ff = 1;
        }
        else
        {
            ff = 0;
        }
        return ff;
    }; // zm--------------------------------------------------------------------------------------------------

    // zm--------------------------------------------------------------------------------------------------
    template <int dim, IVoxNodeType node_type, typename PointType>
    void IVox<dim, node_type, PointType>::IVOX_get_areas(std::vector<float> vec_map_min, std::vector<float> vec_map_max, size_t num_area)
    {
        num_area_ = num_area;
        vec_map_min_ = vec_map_min;
        vec_map_max_ = vec_map_max;
    }; // zm--------------------------------------------------------------------------------------------------

    template <int dim, IVoxNodeType node_type, typename PointType>
    bool IVox<dim, node_type, PointType>::GetClosestPoint(const PointType &pt, PointType &closest_pt)
    {
        std::vector<DistPoint> candidates;
        auto key = Pos2Grid(ToEigen<float, dim>(pt));
        std::for_each(nearby_grids_.begin(), nearby_grids_.end(), [&key, &candidates, &pt, this](const KeyType &delta)
                      {
        auto dkey = key + delta;
        auto iter = grids_map_.find(dkey);
        if (iter != grids_map_.end()) {
            DistPoint dist_point;
            bool found = iter->second->second.NNPoint(pt, dist_point);
            if (found) {
                candidates.emplace_back(dist_point);
            }
        } });

        if (candidates.empty())
        {
            return false;
        }

        auto iter = std::min_element(candidates.begin(), candidates.end());
        closest_pt = iter->Get();
        return true;
    }

    template <int dim, IVoxNodeType node_type, typename PointType>
    bool IVox<dim, node_type, PointType>::GetClosestPoint(const PointType &pt, PointVector &closest_pt, int max_num,
                                                          double max_range)
    {
        std::vector<DistPoint> candidates;
        candidates.reserve(max_num * nearby_grids_.size());

        auto key = Pos2Grid(ToEigen<float, dim>(pt));

// #define INNER_TIMER
#ifdef INNER_TIMER
        static std::unordered_map<std::string, std::vector<int64_t>> stats;
        if (stats.empty())
        {
            stats["knn"] = std::vector<int64_t>();
            stats["nth"] = std::vector<int64_t>();
        }
#endif

        for (const KeyType &delta : nearby_grids_)
        {
            auto dkey = key + delta;
            auto iter = grids_map_.find(dkey);
            if (iter != grids_map_.end())
            {
#ifdef INNER_TIMER
                auto t1 = std::chrono::high_resolution_clock::now();
#endif
                auto tmp = iter->second->second.KNNPointByCondition(candidates, pt, max_num, max_range);
#ifdef INNER_TIMER
                auto t2 = std::chrono::high_resolution_clock::now();
                auto knn = std::chrono::duration_cast<std::chrono::nanoseconds>(t2 - t1).count();
                stats["knn"].emplace_back(knn);
#endif
            }
        }

        if (candidates.empty())
        {
            return false;
        }
#ifdef INNER_TIMER
        auto t1 = std::chrono::high_resolution_clock::now();
#endif

        if (candidates.size() <= max_num)
        {
        }
        else
        {
            std::nth_element(candidates.begin(), candidates.begin() + max_num - 1, candidates.end());
            candidates.resize(max_num);
        }
        std::nth_element(candidates.begin(), candidates.begin(), candidates.end());

#ifdef INNER_TIMER
        auto t2 = std::chrono::high_resolution_clock::now();
        auto nth = std::chrono::duration_cast<std::chrono::nanoseconds>(t2 - t1).count();
        stats["nth"].emplace_back(nth);

        constexpr int STAT_PERIOD = 100000;
        if (!stats["nth"].empty() && stats["nth"].size() % STAT_PERIOD == 0)
        {
            for (auto &it : stats)
            {
                const std::string &key = it.first;
                std::vector<int64_t> &stat = it.second;
                int64_t sum_ = std::accumulate(stat.begin(), stat.end(), 0);
                int64_t num_ = stat.size();
                stat.clear();
                std::cout << "inner_" << key << "(ns): sum=" << sum_ << " num=" << num_ << " ave=" << 1.0 * sum_ / num_
                          << " ave*n=" << 1.0 * sum_ / STAT_PERIOD << std::endl;
            }
        }
#endif

        closest_pt.clear();
        for (auto &it : candidates)
        {
            closest_pt.emplace_back(it.Get());
        }
        return closest_pt.empty() == false;
    }

    template <int dim, IVoxNodeType node_type, typename PointType>
    size_t IVox<dim, node_type, PointType>::NumValidGrids() const
    {
        return grids_map_.size();
    }

    template <int dim, IVoxNodeType node_type, typename PointType>
    void IVox<dim, node_type, PointType>::GenerateNearbyGrids()
    {
        if (options_.nearby_type_ == NearbyType::CENTER)
        {
            nearby_grids_.emplace_back(KeyType::Zero());
        }
        else if (options_.nearby_type_ == NearbyType::NEARBY6)
        {
            nearby_grids_ = {KeyType(0, 0, 0), KeyType(-1, 0, 0), KeyType(1, 0, 0), KeyType(0, 1, 0),
                             KeyType(0, -1, 0), KeyType(0, 0, -1), KeyType(0, 0, 1)};
        }
        else if (options_.nearby_type_ == NearbyType::NEARBY18)
        {
            nearby_grids_ = {KeyType(0, 0, 0), KeyType(-1, 0, 0), KeyType(1, 0, 0), KeyType(0, 1, 0),
                             KeyType(0, -1, 0), KeyType(0, 0, -1), KeyType(0, 0, 1), KeyType(1, 1, 0),
                             KeyType(-1, 1, 0), KeyType(1, -1, 0), KeyType(-1, -1, 0), KeyType(1, 0, 1),
                             KeyType(-1, 0, 1), KeyType(1, 0, -1), KeyType(-1, 0, -1), KeyType(0, 1, 1),
                             KeyType(0, -1, 1), KeyType(0, 1, -1), KeyType(0, -1, -1)};
        }
        else if (options_.nearby_type_ == NearbyType::NEARBY26)
        {
            nearby_grids_ = {KeyType(0, 0, 0), KeyType(-1, 0, 0), KeyType(1, 0, 0), KeyType(0, 1, 0),
                             KeyType(0, -1, 0), KeyType(0, 0, -1), KeyType(0, 0, 1), KeyType(1, 1, 0),
                             KeyType(-1, 1, 0), KeyType(1, -1, 0), KeyType(-1, -1, 0), KeyType(1, 0, 1),
                             KeyType(-1, 0, 1), KeyType(1, 0, -1), KeyType(-1, 0, -1), KeyType(0, 1, 1),
                             KeyType(0, -1, 1), KeyType(0, 1, -1), KeyType(0, -1, -1), KeyType(1, 1, 1),
                             KeyType(-1, 1, 1), KeyType(1, -1, 1), KeyType(1, 1, -1), KeyType(-1, -1, 1),
                             KeyType(-1, 1, -1), KeyType(1, -1, -1), KeyType(-1, -1, -1)};
        }
        else
        {
            // LOG(ERROR) << "Unknown nearby_type!";
        }
    }

    template <int dim, IVoxNodeType node_type, typename PointType>
    bool IVox<dim, node_type, PointType>::GetClosestPoint(const PointVector &cloud, PointVector &closest_cloud)
    {
        std::vector<size_t> index(cloud.size());
        for (int i = 0; i < cloud.size(); ++i)
        {
            index[i] = i;
        }
        closest_cloud.resize(cloud.size());

#pragma omp paralell for
        for (auto i = index.begin(); i != index.end(); i++)
        // std::for_each(std::execution::par_unseq, index.begin(), index.end(), [&cloud, &closest_cloud, this](size_t idx)
        {
            PointType pt;
            if (GetClosestPoint(cloud[*i], pt))
            {
                closest_cloud[*i] = pt;
            }
            else
            {
                closest_cloud[*i] = PointType();
            }
        }
        // );
        return true;
    }

    template <int dim, IVoxNodeType node_type, typename PointType>
    void IVox<dim, node_type, PointType>::AddPoints(const PointVector &points_to_add)
    {
        // std::cout << "grids_map_.size()=" << grids_map_.size() << std::endl;
        std::for_each(points_to_add.begin(), points_to_add.end(), [this](const auto &pt)
                      {
        auto key = Pos2Grid(ToEigen<float, dim>(pt));
        auto iter = grids_map_.find(key);
        if (iter == grids_map_.end()) {
            PointType center;
            center.getVector3fMap() = key.template cast<float>() * options_.resolution_;

            grids_cache_.push_front({key, NodeType(center, options_.resolution_)});
            grids_map_.insert({key, grids_cache_.begin()});

            grids_cache_.front().second.InsertPoint(pt);

            if (grids_map_.size() >= options_.capacity_) {
                printf("capacity_ to big release back!!!\n");
                grids_map_.erase(grids_cache_.back().first);
                grids_cache_.pop_back();
            }
        } else {
            iter->second->second.InsertPoint(pt);
            if(iter->second->second.points_.size()>2){
                iter->second->second.points_.erase(iter->second->second.points_.begin());
            }
            grids_cache_.splice(grids_cache_.begin(), grids_cache_, iter->second);
            grids_map_[key] = grids_cache_.begin();
        } });
    }

    template <int dim, IVoxNodeType node_type, typename PointType>
    void IVox<dim, node_type, PointType>::erasePoints(const PointVector &points_to_erase)
    {
        std::for_each(points_to_erase.begin(), points_to_erase.end(), [this](const auto &pt)
                      {
        auto key = Pos2Grid(ToEigen<float, dim>(pt));
        auto iter = grids_map_.find(key);
        if (iter == grids_map_.end()) {
        } else {
            iter->second->second.points_.clear();
        } });
    }

    // zm--------------------------------------------------------------------------------------------------
    template <int dim, IVoxNodeType node_type, typename PointType>
    void IVox<dim, node_type, PointType>::Mapping_AddPoints(const PointVector &map_points_to_add)
    {
        std::for_each(map_points_to_add.begin(), map_points_to_add.end(), [this](const auto &pt)
                      {
            auto key = Pos2Grid(ToEigen<float, dim>(pt)); // calculate key
            auto iter = grids_map_.find(key);             // find key
            if (iter == grids_map_.end())                 // new key
            {
                PointType center;
                center.getVector3fMap() = key.template cast<float>() * options_.resolution_;
                grids_cache_.push_front({key, NodeType(center, options_.resolution_)});
                grids_map_.insert({key, grids_cache_.begin()});
                grids_cache_.front().second.InsertPoint(pt);
                if (grids_map_.size() >= options_.capacity_)
                {
                //printf("capacity_ to big release back!!!\n");
                    bool ff = true;
                    while (ff)
                    {
                        if (!grids_cache_.empty())
                        {
                            auto ite = grids_cache_.end();
                            --ite;
                            auto key0 = ite->first;
                            std::vector<PointType> pc = ite->second.GetPoints();
                            Eigen::Vector3f c;
                            float x = 0;
                            float y = 0;
                            float z = 0;
                            for (int i = 0; i < pc.size(); i++)
                            {
                                x = x + pc[i].x;
                                y = y + pc[i].y;
                                z = z + pc[i].z;
                            }
                            c[0] = x / pc.size();
                            c[1] = y / pc.size();
                            c[2] = z / pc.size();
                            int a = is_in_area(c);
                            if (a == 1) //在范围内，不可以删
                            {
                                grids_cache_.splice(grids_cache_.begin(), grids_cache_, ite);
                                grids_map_[grids_cache_.back().first] = grids_cache_.begin();
                                ff = true;
                            }
                            else
                            {
                                grids_map_.erase(grids_cache_.back().first);
                                grids_cache_.pop_back();
                                ff = false;
                                break;
                            }
                        }
                        else
                        {
                            ff = false;
                        }
                    }
                }
            }
            else // old key
            {
                iter->second->second.InsertPoint(pt);
                if (iter->second->second.Size() > 2)
                {
                iter->second->second.points_.erase(iter->second->second.points_.begin());
                }
                grids_cache_.splice(grids_cache_.begin(), grids_cache_, iter->second);
                grids_map_[key] = grids_cache_.begin();
            } });
    } // zm--------------------------------------------------------------------------------------------------

    template <int dim, IVoxNodeType node_type, typename PointType>
    Eigen::Matrix<int, dim, 1> IVox<dim, node_type, PointType>::Pos2Grid(const IVox::PtType &pt) const
    {
        return (pt * options_.inv_resolution_).array().round().template cast<int>();
    }

    template <int dim, IVoxNodeType node_type, typename PointType>
    std::vector<float> IVox<dim, node_type, PointType>::StatGridPoints() const
    {
        float num = grids_cache_.size(), valid_num = 0, max = 0, min = 100000000;
        float sum = 0, sum_square = 0;
        for (auto &it : grids_cache_)
        {
            float s = it.second.Size();
            valid_num += s > 0;
            max = s > max ? s : max;
            min = s < min ? s : min;
            sum += s;
            sum_square += s * s;
        }
        float ave = float(sum) / num;
        float stddev = num > 1 ? sqrt((float(sum_square) - num * ave * ave) / (num - 1)) : 0;
        return std::vector<float>{valid_num, ave, max, min, stddev};
    }

    // template <int dim, IVoxNodeType node_type, typename PointType>
    // typename pcl::PointCloud<PointType>::Ptr IVox<dim, node_type, PointType>::GetPointCloud()
    // {
    //     typename pcl::PointCloud<PointType>::Ptr pc = boost::make_shared<pcl::PointCloud<PointType>>();
    //     for (auto &iter : grids_map_)
    //     {
    //         NodeType &node = std::get<1>(*std::get<1>(iter));
    //         std::vector<PointType> &points = node.GetPoints();
    //         pc->points.insert(pc->points.end(), points.begin(), points.end());
    //     }
    //     return pc;
    // }

    // zm--------------------------------------------------------------------------------------------------
    template <int dim, IVoxNodeType node_type, typename PointType>
    typename pcl::PointCloud<PointType>::Ptr IVox<dim, node_type, PointType>::GetPointCloud()
    {
        typename pcl::PointCloud<PointType>::Ptr pc = boost::make_shared<pcl::PointCloud<PointType>>();

        if (std::is_same<NodeType, IVoxNode<PointType, dim>>::value)
        {
            // std::cout << "grids_map_.size()  " << grids_map_.size() << std::endl;
            for (auto &iter : grids_map_)
            {
                NodeType &node = std::get<1>(*std::get<1>(iter));
                for (std::size_t i = 0; i < node.Size(); ++i)
                {
                    PointType point = node.GetPoint(i);
                    pc->points.push_back(point);
                }
            }
            return pc;
        }
        else if (std::is_same<NodeType, IVoxNodePhc<PointType, dim>>::value)
        {
            for (auto &iter : grids_map_)
            {
                NodeType &node = std::get<1>(*std::get<1>(iter));
                for (std::size_t i = 0; i < node.Size(); ++i)
                {
                    PointType point = node.GetPoint(i);
                    pc->points.push_back(point);
                }
            }
            return pc;
        }
    } // zm--------------------------------------------------------------------------------------------------

} // namespace faster_lio

#endif
